-- Author: chleniang@163.com

local _M = {}

local mt = { __index = _M }

local cc_dict = ngx.shared.easy_waf_cc_dict

_M.init = function()
    -- 在此处理一些初始化业务
end

_M.set = function(self, key, value, expire)
    local ok, _ = nil, nil
    if expire > 0 then
        ok, _ = cc_dict:set(key, value, expire)
        if not ok then
            ngx.log(ngx.ERR, "[easy_waf] shared_dict SET failed ")
            return
        end
    else
        ok, _ = cc_dict:set(key, value)
        if not ok then
            ngx.log(ngx.ERR, "[easy_waf] shared_dict SET failed ")
            return
        end
    end
    return ok
end

_M.incr = function(self, key)
    local newval, _ = cc_dict:incr(key, 1)
    if not newval then
        ngx.log(ngx.ERR, "[easy_waf] shared_dict INCR failed ")
        return
    end
    return newval
end

_M.get = function(self, key)
    local res, flag = cc_dict:get(key)
    if not res then
        return
    end
    return res
end

_M.set_ccip = function(self, ip, expire)
    local ipkey = g_ccip_prefix .. ip
    local ok, _ = nil, nil
    if expire > 0 then
        ok, _ = cc_dict:set(ipkey, 1, expire)
        if not ok then
            ngx.log(ngx.ERR, "[easy_waf] shared_dict SET failed ")
            return
        end
    else
        ok, _ = cc_dict:set(ipkey, 1)
        if not ok then
            ngx.log(ngx.ERR, "[easy_waf] shared_dict SET failed ")
            return
        end
    end
    return ok
end

_M.ccip_exist = function(self, ip)
    local ipkey = g_ccip_prefix .. ip
    local res, _ = cc_dict:get(ipkey)
    if not res then
        return false
    end

    return true
end

_M.clean_ccip = function(self, ip)
    local all_keys = cc_dict:get_keys(0)
    if (not ip) or (ip == "all") then
        -- ip = nil 等同于 “all", 清除所有 CC 拦截的ip
        for _, key_val in pairs(all_keys) do
            local from, to, err = ngx.re.find(key_val, g_ccip_prefix, "ijo")
            if from then
                cc_dict:delete(key_val)
            end
        end
    else
        -- 清除指定IP
        local ipkey = g_ccip_prefix .. ip
        for _, key_val in pairs(all_keys) do
            local from, to, err = ngx.re.find(key_val, ipkey, "ijo")
            if from then
                cc_dict:delete(key_val)
                return
            end
        end
    end
end

_M.new = function()
    local self = {}
    return setmetatable(self, mt)
end

return _M
